# Copyright 2021 Ant Group Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
load("@rules_proto_grpc//python:defs.bzl", "python_proto_compile")
load("@rules_python//python:defs.bzl", "py_library")
load("@rules_python//python:packaging.bzl", "py_package")

package(default_visibility = ["//visibility:public"])

exports_files(
    [
        "exported_symbols.lds",
        "version_script.lds",
    ],
    visibility = ["//visibility:private"],
)

pybind_extension(
    name = "libspu",
    srcs = ["libspu.cc"],
    linkopts = select({
        "@bazel_tools//src/conditions:darwin": [
            "-Wl,-exported_symbols_list,$(location //spu:exported_symbols.lds)",
        ],
        "//conditions:default": [
            "-Wl,--version-script,$(location //spu:version_script.lds)",
        ],
    }),
    deps = [
        ":exported_symbols.lds",
        ":version_script.lds",
        "//libspu:version",
        "//libspu/compiler:compile",
        "//libspu/compiler/common:compilation_context",
        "//libspu/core:logging",
        "//libspu/device:api",
        "//libspu/device:io",
        "//libspu/device/pphlo:pphlo_executor",
        "@yacl//yacl/link",
    ],
)

pybind_extension(
    name = "libpsi",
    srcs = ["libpsi.cc"],
    linkopts = select({
        "@bazel_tools//src/conditions:darwin": [
            "-Wl,-exported_symbols_list,$(location //spu:exported_symbols.lds)",
        ],
        "//conditions:default": [
            "-Wl,--version-script,$(location //spu:version_script.lds)",
        ],
    }),
    deps = [
        ":exported_symbols.lds",
        ":version_script.lds",
        "@psi//psi:launch",
        "@psi//psi/legacy:memory_psi",
        "@yacl//yacl/link",
    ],
)

py_library(
    name = "api",
    srcs = [
        "api.py",
        "spu_pb2.py",
        "//libspu:spu_py_proto",
    ],
    data = [
        ":libspu.so",
    ],
)

python_proto_compile(
    name = "psi_py_proto",
    output_mode = "NO_PREFIX_FLAT",
    protos = ["@psi//psi/proto:psi_proto"],
)

python_proto_compile(
    name = "link_py_proto",
    output_mode = "NO_PREFIX_FLAT",
    protos = ["@yacl//yacl/link:link_proto"],
)

python_proto_compile(
    name = "psi_v2_py_proto",
    output_mode = "NO_PREFIX",
    protos = ["@psi//psi/proto:psi_v2_proto"],
)

# Hack generated protobuf due to https://github.com/protocolbuffers/protobuf/issues/1491
genrule(
    name = "psi_v2_py_proto_fixed",
    srcs = [":psi_v2_py_proto"],
    outs = ["psi_v2_pb2.py"],
    cmd = "sed 's#from yacl.link import#from . import#g;s#from psi.proto import#from . import#g' $(SRCS) > $(OUTS)",
)

py_library(
    name = "psi",
    srcs = [
        "psi.py",
        ":link_py_proto",
        ":pir_py_proto",
        ":psi_py_proto",
        ":psi_v2_py_proto_fixed",
    ],
    data = [
        ":libpsi.so",
        ":libspu.so",
    ],
)

python_proto_compile(
    name = "pir_py_proto",
    output_mode = "NO_PREFIX_FLAT",
    protos = ["@psi//psi/proto:pir_proto"],
)

py_library(
    name = "init",
    srcs = [
        "__init__.py",
        "version.py",
        ":api",
        ":psi",
        "//spu/experimental:experimentals",
        "//spu/intrinsic:all_intrinsics",
        "//spu/utils:simulation",
    ],
)

py_package(
    name = "_lib_deps",
    deps = [
        ":api",
    ],
)
